# -*- coding: utf-8 -*-
from pymongo import UpdateOne
from zc_core.client.mongo_client import Mongo

mongo = Mongo()
sp_list = mongo.list('supplier_pool')

full_to_short = dict()
name_to_id = dict()
fullname_to_id = dict()
for sp in sp_list:
    id = sp.get('_id')
    name = sp.get('name')
    full_name = sp.get('fullName')
    full_to_short[full_name] = name
    name_to_id[name] = id
    fullname_to_id[full_name] = id

print('full_to_short: %s' % len(full_to_short))
print('name_to_id: %s' % len(name_to_id))
print('fullname_to_id: %s' % len(fullname_to_id))


def update_supplier(coll):
    print('开始：%s' % coll)
    mongo = Mongo()
    bulk_list = list()
    items = mongo.list(coll, fields={'_id': 1, 'supplierId': 1, 'supplierName': 1})
    for item in items:
        sp_name = item.get('supplierName')
        sp_id = fullname_to_id.get(sp_name)
        if not sp_id:
            sp_id = name_to_id.get(sp_name)
        if not sp_id:
            sp_id = name_to_id.get(sp_name)
        short_name = full_to_short.get(sp_name)
        if short_name:
            item['supplierName'] = short_name
        if sp_id:
            item['supplierId'] = sp_id
            bulk_list.append(UpdateOne({'_id': item.get('_id')}, {'$set': item}, upsert=False))

        if len(bulk_list) > 30000:
            mongo.bulk_write(coll, bulk_list)
            print('更新1：%s' % len(bulk_list))
            bulk_list.clear()
            bulk_list = list()

    if bulk_list:
        print('更新2：%s' % len(bulk_list))
        mongo.bulk_write(coll, bulk_list)
    print('成功：%s' % coll)


if __name__ == '__main__':
    batch_list = '20200913', '20200906', '20200830', '20200823', '20200816', '20200809', '20200802', '20200726', '20200719', '20200712', '20200705', '20200628', '20200621', '20200614', '20200607', '20200531', '20200524', '20200517', '20200510', '20200503', '20200426', '20200419', '20200412', '20200405', '20200329', '20200322', '20200315', '20200308', '20200301', '20200223', '20200216', '20200209', '20200119', '20200112', '20200105'
    for batch_no in batch_list:
        update_supplier('data_{}'.format(batch_no))
